# -*- coding: utf-8 -*-

import grpc
import example_pb2
import example_pb2_grpc 
import os
from PIL import Image
import torchvision.transforms as transforms
import glob
import matplotlib.pyplot as plt
import numpy as np
import time
import torch

preprocessing_train = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
  transforms.Normalize(mean=[0.485, 0.456, 0.406],
                       std=[0.229, 0.224, 0.225]),
])

preprocessing_val = transforms.Compose([
    transforms.Resize(224),
    # transforms.CenterCrop(224),
    # transforms.ToTensor(),
    # transforms.Normalize(mean=[0.485, 0.456, 0.406],
    #                      std=[0.229, 0.224, 0.225]),
])


def run(images, gt_label):
  # with grpc.insecure_channel('localhost:50051') as channel:
  with grpc.insecure_channel('175.9.192.228:50051') as channel:
    stub = example_pb2_grpc.ResNetStub(channel)
    request = example_pb2.ImageMatrix()
    # request.image_matrix.extend([10]*1*3*224*224)
    request.image_matrix.extend(images)
    start = time.time()
    response = stub.ClassifyImage(request)
    end = time.time()
  print("image category: " + str(int(response.category)),
        "  name: ", id_to_label[int(response.category)],
        "  gt_category: ", id_to_label[gt_label],
        "  running time : ", end-start)

if __name__ == '__main__':
  a = os.listdir("/media/retoo/RetooDisk/wanghui/Data/ILSVR2012C/ILSVRC2012_img_val/")
  category = open("../imagenet1000_clsidx_to_labels.txt","r").readlines()
  id_to_label = dict()
  for line in category:
    line = line.split(":")
    id_to_label[int(line[0])] = line[1].strip()

  root = "/media/retoo/RetooDisk/wanghui/Data/ILSVR2012C/ILSVRC2012_img_val/"
  files = open("/media/retoo/RetooDisk/wanghui/Data/ILSVR2012C/ILSVRC2012_val.txt", "r")
  for line in files.readlines():
    line = line.strip().split(" ")
    gt_label = int(line[1])
    p = os.path.join(root, line[0].split("/")[-1])
    src = Image.open(p).convert('RGB')
    # img = preprocessing_val(src)
    src = src.resize((224,224))

    src = np.asarray(src)
    img = torch.from_numpy(src).permute((2, 0, 1)).contiguous()
    img = img.unsqueeze(dim=0)
    run(img.view(-1),gt_label)
    plt.imshow(np.asarray(src))
    plt.show()


# docker build --pull -t resnet-libtorch-serving -f libtorch_cpu_Dockerfile .
# docker run -p 175.9.192.228:50051:50051 --name=resnet_service -d -it resnet-libtorch-serving:latest /bin/bash -c './resnet_server'
# docker run -p 150051:50051 --name=resnet_service -d -it resnet-libtorch-serving:latest /bin/bash -c './resnet_server'
